# @Time    : 2023/11/07 22:52
# @Author  : yaodong
# @Email   : yaodong@lingxi.ai
# @File    : planner_seat.py
# @Software: LLM_internal

# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
import asyncio
import json
import traceback
import uuid

from bot.insurance_planner_gpt.utils import date_util
from bot.insurance_planner_gpt.utils.seat_planner_msg_channel import SeatPlannerMsgChannel
from common import constants
from common.log import logger
from common.message import MessageContent, Message, ConversationMessage
from db.redis.redis_service import RedisMemory
from util_tool.conversation_util import ConversationUtil

cu = ConversationUtil()
redis_memory = RedisMemory()
loop = asyncio.get_event_loop()


class SeatProcessPlannerMsgChannel(SeatPlannerMsgChannel):

    def __init__(self):
        super().__init__()

    async def handle_planner_channel_data(self, session_id, clients):
        """
        坐席工作台 接收用户的消息
        """
        pubsub = self.aio_redis.pubsub()
        await pubsub.subscribe(self.planner_channel)
        try:
            # 循环监听消息
            async for ori_message in pubsub.listen():
                if ori_message['type'] == 'message':
                    message = ori_message['data'].decode('utf-8')
                    context_message_all = ConversationMessage.from_json(message)
                    user_session_id = context_message_all.session_id
                    messages = context_message_all.messages
                    client_type = context_message_all.client_type
                    client_id = context_message_all.client_id
                    action = context_message_all.action
                    last_message = messages[-1]
                    role = last_message.get('role', None)
                    message_id = last_message.get('message_id', str(uuid.uuid4()).replace('-', ''))
                    message_id_process = redis_memory.hget("planner_channel:message_id", message_id)
                    if not message_id_process:
                        redis_memory.hset("planner_channel:message_id", message_id, 1)
                        redis_memory.expireat("planner_channel:message_id", date_util.get_next_day_start())
                        last_content_data = last_message['content'][0].data
                        msg_type = last_message['content'][0].msg_type
                        logger.info(
                            f"会话id:{session_id},从{self.planner_channel}中获取到一条用户消息:{context_message_all.to_json()}")
                        agent_client_ids = redis_memory.smembers("agent_connect:" + user_session_id)
                        # 处理接管情况下的规划师响应
                        if action == "takeover_response":
                            takeover_client_id = redis_memory.get("agent_takeover:" + user_session_id)
                            for agent_connect_client_id in agent_client_ids:
                                if agent_connect_client_id != takeover_client_id:
                                    msg_start = ConversationMessage.gen_msg(constants.MESSAGE_TYPE_TEXT, client_id,
                                                                            client_type,
                                                                            "▌", user_session_id, user_session_id,
                                                                            code=constants.MESSAGE_CODE_END,
                                                                            role=role, message_id=message_id)
                                    common_msg = ConversationMessage.gen_msg(constants.MESSAGE_TYPE_TEXT, client_id,
                                                                             client_type,
                                                                             last_content_data, user_session_id,
                                                                             user_session_id,
                                                                             code=constants.MESSAGE_CODE_END, role=role,
                                                                             message_id=message_id)
                                    await send_to_seat([agent_connect_client_id], msg_start, clients)
                                    await send_to_seat([agent_connect_client_id], common_msg, clients)
                            takeover_msg = ConversationMessage.gen_msg(constants.MESSAGE_TYPE_EDIT_TEXT, client_id,
                                                                       client_type, last_content_data, user_session_id,
                                                                       user_session_id,
                                                                       code=constants.MESSAGE_CODE_END,
                                                                       role=role, message_id=message_id)
                            await send_to_seat([takeover_client_id], takeover_msg, clients)
                        elif action == "agent_solution":
                            common_msg = ConversationMessage.gen_msg(msg_type, client_id, client_type,
                                                                     last_content_data,
                                                                     user_session_id, user_session_id,
                                                                     code=constants.MESSAGE_CODE_END,
                                                                     role=role,
                                                                     message_id=message_id)
                            await send_to_seat(agent_client_ids, common_msg, clients)
                        else:
                            # 处理用户的对话
                            common_msg = ConversationMessage.gen_msg(msg_type, client_id,
                                                                     client_type,
                                                                     last_content_data,
                                                                     user_session_id, user_session_id,
                                                                     code=constants.MESSAGE_CODE_END,
                                                                     role=role,
                                                                     message_id=message_id)
                            await send_to_seat(agent_client_ids, common_msg, clients)
        except Exception:
            logger.error(f"处理{self.planner_channel}通道消息异常: {traceback.format_exc()}")
        finally:
            await pubsub.close()
            await pubsub.unsubscribe(self.planner_channel)


class PlannerSeat:

    def __init__(self, websocket, client_id, clients, client_type="agent_work"):
        self.cu = ConversationUtil()
        self.websocket = websocket
        self.clients = clients
        self.client_id = client_id
        self.client_type = client_type
        self.planner_channel = SeatProcessPlannerMsgChannel()

    async def handle_event(self, agent_action, user_session_id, agent_client_id, conversation, messages):
        """ 统一处理坐席事件"""
        try:
            if agent_action == "agent_connect":
                await self.handle_agent_connect(user_session_id, agent_client_id)
            elif agent_action == "agent_takeover":
                await self.handle_agent_takeover(user_session_id, agent_client_id, self.client_type)
            elif agent_action == "agent_stop_takeover":
                await self.handle_agent_stop_takeover(user_session_id, agent_client_id,
                                                      self.client_type)
            elif agent_action == "agent_transfer":
                await self.handle_agent_transfer(user_session_id, self.client_type, messages)
            elif agent_action == "agent_solution_update":
                await self.handle_agent_solution_update(user_session_id, messages)
            elif agent_action == "agent_login":
                await self.handle_agent_login(agent_client_id)
            elif agent_action == "agent_logout":
                await self.handle_agent_log_out(agent_client_id)
        except Exception:
            logger.error("处理坐席消息异常" + traceback.format_exc())

    async def handle_agent_login(self, client_id):
        """ 坐席登录 """
        try:
            self.clients[client_id] = self.websocket
        except Exception:
            pass

    async def handle_agent_log_out(self, client_id):
        """ 坐席登出 """
        try:
            # 删除退出坐席的链接
            del self.clients[client_id]
            # 清除所有坐席链接
            redis_memory.find_and_remove_value(key_prefix="agent_connect", value_to_remove=client_id)

            keys_to_delete = redis_memory.keys('agent_takeover:*')
            for key in keys_to_delete:
                if redis_memory.get(key) == client_id:
                    redis_memory.delete(key)
        except Exception:
            pass

    async def handle_agent_connect(self, user_session_id, agent_client_id):
        """ 坐席主动连接一个用户会话 """
        try:
            redis_memory.sadd("agent_connect:" + user_session_id, agent_client_id)

            # 返回连接成功
            await self.send_msg_to_seat(constants.MESSAGE_TYPE_TEXT, self.client_type, agent_client_id,
                                        "连接成功", constants.MESSAGE_CODE_SUCCESS, None, user_session_id)
        except Exception:
            pass

    async def handle_agent_takeover(self, user_session_id, agent_client_id, client_type):
        """ 坐席主动接管一个用户会话 """
        try:
            agent_client_id_exist = redis_memory.get("agent_takeover:" + user_session_id)
            if agent_client_id_exist:
                # 如果已经有坐席在接管了 则返回接管失败
                await self.send_msg_to_seat(constants.MESSAGE_TYPE_TEXT, client_type, agent_client_id,
                                            "接管失败", constants.MESSAGE_CODE_SUCCESS, None,
                                            agent_client_id)
                return
            redis_memory.set("agent_takeover:" + user_session_id, agent_client_id, ex=3600)
            # 返回接管成功
            await self.send_msg_to_seat(constants.MESSAGE_TYPE_TEXT, client_type, agent_client_id,
                                        "接管成功", constants.MESSAGE_CODE_SUCCESS, None, agent_client_id)

        except Exception:
            logger.info(f"接管异常:{traceback.format_exc()}")

    async def handle_agent_stop_takeover(self, user_session_id, agent_client_id, client_type):
        """ 坐席停止接管一个用户会话 """
        try:
            redis_memory.delete("agent_takeover:" + user_session_id)
            # 返回停止接管成功
            await self.send_msg_to_seat(constants.MESSAGE_TYPE_TEXT, client_type, user_session_id,
                                        "停止接管", constants.MESSAGE_CODE_SUCCESS, None, agent_client_id)

        except Exception:
            logger.info(f"异常:{traceback.format_exc()}")

    async def handle_agent_transfer(self, user_session_id, client_type, messages):
        """ 坐席发送消息给用户 """
        try:
            if user_session_id:

                message_id_temp = messages[0]['message_id']
                agent_update_message = messages[0]['content'][0].data

                # 给所有正在连接的坐席发送规划师响应 丢到队列
                msg = ConversationMessage.gen_msg(constants.MESSAGE_TYPE_TEXT, "", client_type,
                                                  agent_update_message,
                                                  user_session_id, user_session_id, role=constants.ROLE_ASSISTANT)

                await self.planner_channel.publish_msg(self.planner_channel.seat_channel, msg)
                # 存储坐席修改对话
                if message_id_temp:
                    loop.run_in_executor(None, cu.store_conversation_index, user_session_id,
                                         "agent_" + message_id_temp,
                                         constants.MESSAGE_TYPE_TEXT, agent_update_message,
                                         constants.ROLE_ASSISTANT, self.client_type)
                redis_memory.expire("agent_takeover:" + user_session_id, 3600)

        except Exception:
            logger.info(f"异常:{traceback.format_exc()}")

    async def handle_agent_solution_update(self, user_session_id, messages):
        """ 坐席修改规划师AGENT推断数据 """
        if user_session_id and user_session_id in self.clients:

            data = messages[0]['content'][0].data
            data_dict = json.loads(data)
            # 获取 'title' 和 'content' 字段
            title = data_dict.get('title', '')  # 如果没有 'title' 字段，返回默认值
            content = data_dict.get('content', '')  # 如果没有 'content' 字段，返回默认值

            if '潜在需求' == title:
                if redis_memory.exist_key("planner:task:" + user_session_id):
                    redis_memory.set_memory("planner:task:" + user_session_id, "PlanAssume", content)
            elif '用户问题' == title:
                if redis_memory.exist_key("planner:task:" + user_session_id):
                    redis_memory.set_memory("planner:task:" + user_session_id, "QuestioningDisputeResolution",
                                            content)
            elif '因果推断' == title:
                if redis_memory.exist_key("planner:task:" + user_session_id):
                    redis_memory.set_memory("planner:task:" + user_session_id, "ProductCausalSolution", content)

            # 存储坐席修改后的AGENT方案信息
            message_id_temp = messages[0]['message_id']
            if message_id_temp:
                loop.run_in_executor(None, cu.store_conversation_index, user_session_id,
                                     "agent_" + message_id_temp,
                                     constants.MESSAGE_TYPE_SOLUTION, messages[0]['content'][0].data,
                                     constants.ROLE_ASSISTANT, self.client_type)

            redis_memory.expire("agent_takeover:" + user_session_id, 3600)

    async def send_msg_to_seat(self, msg_type, client_type, user_session_id, message, code, message_id,
                               agent_client_id="", role=constants.ROLE_AGENT):
        """ 发送消息给坐席平台 """
        try:
            if self.clients and agent_client_id and agent_client_id in self.clients:
                websocket = self.clients[agent_client_id]
                if websocket:
                    # 构造消息格式
                    send_content = MessageContent(msg_type, message)
                    content_arr = [send_content]
                    send_message = Message(message_id, role, content_arr, None)
                    send_msg = ConversationMessage(client_type=client_type, client_id=agent_client_id,
                                                   session_id=agent_client_id,
                                                   user_name='',
                                                   target_client_id=agent_client_id, messages=send_message, code=code)
                    await websocket.send(send_msg.to_json())

        except Exception:
            logger.info(f"用户会话id:{user_session_id},给坐席发送对话异常,{traceback.format_exc()}")


async def send_to_seat(agent_client_ids, message, clients):
    """ 发送消息给坐席平台 """
    try:
        if agent_client_ids:
            for agent_client_id in agent_client_ids:
                if clients and agent_client_id and agent_client_id in clients:
                    websocket = clients[agent_client_id]
                    if websocket:
                        await websocket.send(message.to_json())
    except Exception:
        logger.error(f"发送坐席平台异常")
